import os
import tempfile
import paramiko

class Connection(object):
    """Connects and logs into the specified hostname. 
    Arguments that are not given are guessed from the environment.""" 

    
    
    def sshKeyShare(self):
        
        
        print "---------------------------------------------------------------"
        print "Setting up Passwordless SSH"  
        print "---------------------------------------------------------------"
        nodeKeyFiles = []
        knownHostsString = ""
           
        for node in self.nodes:
        
            sshCon = ssh.Connection(node['nodeName'], username=self.hadoopUser, password=self.hadoopPW)
            remotePath = "/home/" + self.hadoopUser + "/.ssh"
            keyGenString = "ssh-keygen -t rsa -P '' -f " + remotePath + "/id_rsa  < <(echo y) "
            renameString = "cat " + remotePath + "/id_rsa.pub >> " + remotePath + "/authorized_keys"
            nodeIP = socket.gethostbyname(node['nodeName'])
            knownHostsString = knownHostsString + node['nodeName'] + "," + nodeIP + " "
            sshCon.execute(keyGenString)
            sshCon.execute(renameString)
            remoteName = "./conf/ssh/" + node['nodeName'] + "-authorized_keys"
            sshCon.get(remotePath + "/authorized_keys", remoteName)
            nodeKeyFiles.append(remoteName)
            sshCon.close()
        
        fullKeyFile = open("./conf/ssh/authorized_keys", "w")
        
        for nodeKeyFile in nodeKeyFiles:
            nkFile = open(nodeKeyFile, "r")
            fullKeyFile.write(nkFile.read())
       
        fullKeyFile.close()
        nkFile.close()
        
        
        knownHosts = open("./conf/ssh/known_hosts", "w")
        sshCon = ssh.Connection(self.nodes[0]['nodeName'], username=self.hadoopUser, password=self.hadoopPW)
        knownHostsCMD = "ssh-keyscan -t rsa " + knownHostsString
        knownHostResults = sshCon.execute(knownHostsCMD)
        for knownHost in knownHostResults:
            knownHosts.write(knownHost + "\n")
        knownHosts.close()

        
        
            
        for node in self.nodes:
            sshCon = ssh.Connection(node['nodeName'], username=self.hadoopUser, password=self.hadoopPW)
            sshCon.execute("mkdir -p " + remotePath)
            sshCon.put("./conf/ssh/authorized_keys", remotePath + "/authorized_keys")
            sshCon.put("./conf/ssh/known_hosts", remotePath + "/known_hosts")
            
            sshCon.execute("chmod 0600 " + remotePath + "/authorized_keys")
            sshCon.execute("chmod 0700 " + remotePath)
            print "    " + node['nodeName'] + " : " + "authorized_keys placed in /home/" + self.hadoopUser + "/.ssh"
   

         
    
    
    
    def __init__(self,
                 host,
                 username = None,
                 private_key = None,
                 password = None,
                 port = 22,
                 ):
        self._sftp_live = False
        self._sftp = None
        if not username:
            username = os.environ['LOGNAME']

        # Log to a temporary file.
        templog = tempfile.mkstemp('.txt', 'ssh-')[1]
        paramiko.util.log_to_file(templog)

        # Begin the SSH transport.
        self._transport = paramiko.Transport((host, port))
        self._tranport_live = True
        # Authenticate the transport.
        if password:
            # Using Password.
            self._transport.connect(username = username, password = password)
        else:
            # Use Private Key.
            if not private_key:
                # Try to use default key.
                if os.path.exists(os.path.expanduser('~/.ssh/id_rsa')):
                    private_key = '~/.ssh/id_rsa'
                elif os.path.exists(os.path.expanduser('~/.ssh/id_dsa')):
                    private_key = '~/.ssh/id_dsa'
                else:
                    raise TypeError, "You have not specified a password or key."

            private_key_file = os.path.expanduser(private_key)
            rsa_key = paramiko.RSAKey.from_private_key_file(private_key_file)
            self._transport.connect(username = username, pkey = rsa_key)
    
    def _sftp_connect(self):
        """Establish the SFTP connection."""
        if not self._sftp_live:
            self._sftp = paramiko.SFTPClient.from_transport(self._transport)
            self._sftp_live = True

    def get(self, remotepath, localpath = None):
        """Copies a file between the remote host and the local host."""
        if not localpath:
            localpath = os.path.split(remotepath)[1]
        self._sftp_connect()
        self._sftp.get(remotepath, localpath)

    def put(self, localpath, remotepath = None):
        """Copies a file between the local host and the remote host."""
        if not remotepath:
            remotepath = os.path.split(localpath)[1]
        self._sftp_connect()
        self._sftp.put(localpath, remotepath)

    def execute(self, command):
        """Execute the given commands on a remote machine."""
        channel = self._transport.open_session()
        channel.exec_command(command)
        output = channel.makefile('rb', -1).readlines()
        if output:
            return output
        else:
            return channel.makefile_stderr('rb', -1).readlines()

    def executeSt(self, command):
        """Execute the given commands on a remote machine."""
        channel = self._transport.open_session()
        channel.exec_command(command)
        exitStatus = channel.recv_exit_status()
        output = channel.makefile('rb', -1).readlines()
        if output:
            return (exitStatus,output)
        else:
            return (exitStatus,channel.makefile_stderr('rb', -1).readlines())




    def close(self):
        """Closes the connection and cleans up."""
        # Close SFTP Connection.
        if self._sftp_live:
            self._sftp.close()
            self._sftp_live = False
        # Close the SSH Transport.
        if self._tranport_live:
            self._transport.close()
            self._tranport_live = False

    def __del__(self):
        """Attempt to clean up if not explicitly closed."""
        self.close()

def main():
    """Little test when called directly."""
    # Set these to your own details.
    myssh = Connection('example.com')
    myssh.put('ssh.py')
    myssh.close()

# start the ball rolling.
if __name__ == "__main__":
    main()
